
import numpy as np
import os
import mindspore.dataset as ds
from PIL import Image
import cv2
import random
import glob
from mindvideo.common.utils.class_factory import ClassFactory, ModuleType
import mindspore.dataset.vision.c_transforms as c_vision
@ClassFactory.register(ModuleType.DATASET)
class UCFDataset:
    def __init__(self, dataset_dir, seq_len=16, label_path=None):
        self.filelist=self.get_filelist(dataset_dir)
        if label_path:
            self.label_dict=self.load_label(label_path)
        else:
            self.label_dict=self.create_label(self.filelist)
            
        self.seq_len = seq_len
    def __getitem__(self, index):
        path=self.filelist[index]
        self.data=self.load_video(path)
        # print(path, self.data.shape)
        class_name=path.split('_')[-3]
        return self.data, self.label_dict[class_name]

    def __len__(self):
        return len(self.filelist)
    
    def get_filelist(self, sample_path): 
        # filelist = []
        files = os.listdir(sample_path)
        filelist = [os.path.join(sample_path, f) for f in files]
        # for home, dirs, _ in os.walk(sample_path):
        #     if home == sample_path:
        #         continue
        #     for dir in dirs:
        #         filelist.append(os.path.join(home,dir))
        return filelist

    def load_video(self, video_path):
        data=[]
        idx=0
        path_file_number=glob.glob(os.path.join(video_path, '*.jpg'))#文件个数

        video_len=len(path_file_number)
        clipnum=int(video_len/self.seq_len)
        # filename=video_path.split('/')[-1]
        if clipnum==1:
            idx==0
        else:
            idx=np.random.randint(0,clipnum)
        for i in range (self.seq_len*idx, self.seq_len*idx+self.seq_len):
            # img = Image.open(path+'/'+filename+"_"'{}.jpg'.format(i))	#读取文件
            img = Image.open("{}/{}.jpg".format(video_path, i))
            img = np.array(img, dtype = np.float64)	#转为float64类型的Numpy数组
            data.append(img)
        return np.array(data)

    def load_label(self, label_path):     
        label={}
        with open(label_path, 'r')as f:
            while True:
                line=f.readline()
                if not line:
                    break
                # 跳过第一行
                if line.split(' ')[0] == "Class_id":
                    continue
                label[line.split()[-1]]=int(line.split()[0])
        return label
    def create_label(self, filelist):
        files = sorted(filelist)
        cls_cnt = 0
        label_dict = {}
        for file_name in files:
            cls_name = file_name.split('_')[-3]
            if cls_name not in label_dict.keys():
                label_dict[cls_name] = cls_cnt
                cls_cnt+=1
        return label_dict

# from mindvideo.datasets.preprocess import ReOrder, VideoResize
# dataset_generator = UCFDataset(dataset_dir="mindvideo/datasets/UCF101/data/train", 
#                             #    label_path="mindvideo/datasets/UCF101/labels.txt",
#                                seq_len=64)
# transforms_list = [VideoResize([256, 256]),
#                    ReOrder([3,0,1,2])]
# dataset = ds.GeneratorDataset(dataset_generator, ["data", "label"], shuffle=False)
# dataset = dataset.map(operations=transforms_list, input_columns="data")
# dataset = dataset.batch(batch_size=16, drop_remainder=True)

# for data in dataset.create_dict_iterator():
#     print("video shape: {}".format(data['data'].shape), ", Label: {}".format(data['label']))
